import pickle
import umap
from ..config import (
    gaobao_data_dir,
    dmeta_embedding_dir,
    embed_save_dir,
    context_size,
    chunk_size,
)
from ..chunk.CharacterChunker import text_splitter
from langchain_community.document_loaders import JSONLoader
from langchain_core.documents import Document
from ..sent2vec import Sent2VecEmbeddings
from .BIC import bic
import torch
import os
import json
from ..chunk.CharacterChunker import text_splitter


def metadata_func(record: dict, metadata: dict) -> dict:

    metadata["title"] = record.get("title")
    metadata["dir"] = record.get("dir")

    # if "source" in metadata:
    #     source = metadata["source"].split("/")
    #     source = source[source.index("langchain"):]
    #     metadata["source"] = "/".join(source)

    return metadata


def collate(numbers, max_number):
    # 给定一个数字列表，可以按照相同的元素把他们聚集起来获得一个二维数组。
    # 比如二维列表的第0号位置是数字列表里是0的元素的所有下标。
    two_dimensional_list = [
        [] for _ in range(max_number + 1)
    ]  # 创建一个二维列表，行数为最大值加一

    for index, value in enumerate(numbers):
        two_dimensional_list[value].append(
            index
        )  # 根据数字的值将下标添加到对应数字的列表中

    return two_dimensional_list

def special_collate(labels,idx):
    assert len(labels) == len(idx),"因为labels是根据idx预测的，理论上不可能不同"
    # 很难形容，专用于拆分过大内容用的
    max_label=max(labels)
    
    two_dimensional_list = [
        [] for _ in range(max_label + 1)
    ]  # 创建一个二维列表，行数为最大值加一

    for index, value in enumerate(labels):
        two_dimensional_list[value].append(
            idx[index]
        )  # 根据数字的值将下标添加到对应数字的列表中

    return two_dimensional_list

loader = JSONLoader(
    file_path=gaobao_data_dir,
    jq_schema=".[]",
    text_content=True,
    content_key="content",
    metadata_func=metadata_func,
)

sent2vec = Sent2VecEmbeddings(dmeta_embedding_dir)

# 要从第几层开始,所有的逻辑分为三类，一类是可以复用，一类是从头开始，一类是不能复用
layer_cnt = 0
os.makedirs(f"{embed_save_dir}/{layer_cnt}", exist_ok=True)

if layer_cnt>0 and os.path.exists(f"{embed_save_dir}/{layer_cnt-1}/current_texts.json"):
    current_texts=json.load(open(f"{embed_save_dir}/{layer_cnt-1}/current_texts.json"))
elif layer_cnt>0 and not os.path.exists(f"{embed_save_dir}/{layer_cnt-1}/current_texts.json"):
    print('因为不能获取历史层的所有文本信息，不能跳层')
    exit()
else:
    # 进这里肯定是0层
    assert layer_cnt==0
    
    # BUG 注意，这里为了调试只开了100
    data = loader.load()[:100]
    data = text_splitter.transform_documents(data)
    texts = [d.page_content for d in data]
    print(f'分块后的texts有{len(texts)}个')
    # 太短的文本就不要了
    for  i in range(len(texts)-1,-1,-1):
        if len(texts[i])<10:
            del texts[i]
    with open(
            f"{embed_save_dir}/{layer_cnt}/current_texts.json", "w", encoding="utf-8"
        ) as o:
            # pickle似乎不能处理二维的数组
            json.dump(texts, o,ensure_ascii=False, indent=4)

    # 这个过程有点像层次遍历其实
    current_texts = texts

if layer_cnt>0 and os.path.exists(f"{embed_save_dir}/{layer_cnt-1}/text_need_embed"):
    with open(f"{embed_save_dir}/{layer_cnt}/text_need_embed", "wb") as o:
            text_need_embed=pickle.load(o)
elif layer_cnt>0 and not os.path.exists(f"{embed_save_dir}/{layer_cnt-1}/text_need_embed"):
    print('因为不能获取当前层应该处理的数据，不能跳层')
    exit()
else:
    text_need_embed = (
        current_texts  # 为了每次新的一轮不需要重复上一轮的编码过程，只编码新增的元素设立
    )



if os.path.exists(f"{embed_save_dir}/{layer_cnt}/record"):
    record=pickle.load(open(f"{embed_save_dir}/{layer_cnt-1}/record",'rb'))
elif layer_cnt==0:
    record = set()
else:
    print('如果没有保存过历史重复的记录，是不允许跳层继续迭代的。')
    exit()
    
while True:

    # 如果是为了复用，我们应该存的是当前层的embedding，同时编码当前层所需要的东西
    # 如果之前就没有保存过，是不可能直接跨层的，因为record不匹配。
    if layer_cnt>0 and os.path.exists(f"{embed_save_dir}/{layer_cnt-1}/embedding.pt"):
        embedding = torch.load(f"{embed_save_dir}/{layer_cnt-1}/embedding.pt")
        current_embedding = sent2vec.embed_documents(text_need_embed)[0]
        embedding = torch.cat((embedding, current_embedding), dim=0)
    elif layer_cnt>0 and not os.path.exists(f"{embed_save_dir}/{layer_cnt-1}/embedding.pt"):
        exit()
    elif layer_cnt==0 and os.path.exists(f"{embed_save_dir}/{0}/embedding.pt"):
        embedding = torch.load(f"{embed_save_dir}/{0}/embedding.pt")
    else:
        embedding = sent2vec.embed_documents(text_need_embed)[0]
    torch.save(embedding, f"{embed_save_dir}/{layer_cnt}/embedding.pt")


    labels = bic(embedding)
    max_labels = max(labels)
    if max_labels == 0:
        break  # 说明没有不同的聚类了
    indices = collate(labels, max_number=max_labels)

    next_text = []
    record_cluster_text = []  # just for debug use
    # 先递归的处理,但是这里会不会有死循环的可能呢？
    del_idx=[] # 这里面的索引都是需要被拆解的，拆解完成后就不需要他了。
    # print(indices)
    for i, idx in enumerate(indices):
        if len(idx) > context_size//chunk_size:
            # 先取出这组idx对应的embedding,然后取出label
            temp_label = bic(torch.index_select(embedding,0,torch.tensor(idx)))
            # 这里就麻烦了，因为collate是根据元素索引位置来确定idx的,但是现在idx实际上要去idx里面取
            temp_indices = special_collate(temp_label, idx=idx)
            del_idx.append(i)
            indices.extend(temp_indices)
    # print(f'删除之前的簇数有{len(indices)}')
    
    # print('超出的数量是',context_size//chunk_size)
    for i,d in enumerate(del_idx):
        print(i,indices[d])
    
    for index in del_idx:
        del indices[index]
    
    # print(f'最后处理出来的簇数有{len(indices)}')
    # print(indices)
    for i, idx in enumerate(indices):
        # 如果这个聚类是见过的，那就不整了。
        if tuple(idx) not in record:
            record.add(tuple(idx))
        else:
            continue


        cluster_text = [current_texts[i] for i in idx]
        record_cluster_text.append(cluster_text)
        # 如果idx里只有一个元素，那么就是没人和他一个类，不用摘要,这种情况下什么都不需要做
        # if len(idx)>=2:
        # TODO
        # summary_text=summary_generate(cluster_text)
        # summary_text=text_splitter.split_text(summary_text)
        # next_text.append(summary_text)

        # # 二次切分next_text，防止超chunk。

    layer_cnt += 1
    os.makedirs(f"{embed_save_dir}/{layer_cnt}", exist_ok=True)
    with open(f"{embed_save_dir}/{layer_cnt}/content.json", "w", encoding="utf-8") as o:
        json.dump(
            { "content": next_text},
            o,
            ensure_ascii=False,
            indent=4,
        )
    with open(
        f"{embed_save_dir}/{layer_cnt}/debug_cluster.json", "w", encoding="utf-8"
    ) as o:
        json.dump(record_cluster_text, o, ensure_ascii=False, indent=4)
    with open(f"{embed_save_dir}/{layer_cnt}/record", "wb") as o:
        pickle.dump(record, o)
    
    # 保存了当前层及所有子层全部的文本
    with open(
        f"{embed_save_dir}/{layer_cnt}/current_texts.json", "w", encoding="utf-8"
    ) as o:
        json.dump(current_texts, o,ensure_ascii=False, indent=4)
    exit()
    # 把二维切分的next_text展开，因为current text不是二维的。
    flattened_next_text = [element for sublist in next_text for element in sublist]

    # 如果某一轮都不能生成摘要了，说明生成结束了？
    if len(flattened_next_text) == 0:
        print(f'因为flattened_next_text长度是0退出')
        break
    text_need_embed = flattened_next_text
    with open(f"{embed_save_dir}/{layer_cnt}/text_need_embed", "wb") as o:
        pickle.dump(text_need_embed,o)
    current_texts.extend(flattened_next_text)
